import torch
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint

from mec.data import MulEncDataset
from mec.model import MulEncNet

num_gpu = 1 if torch.cuda.is_available() else 0


def MuEncTrainTest(md: MulEncDataset,
                   batch_size: int = 32,
                   max_out: int = 20,
                   layers: int = 1,
                   embedding_size: int = 128,
                   hidden_size: int = 256,
                   patience: int = 10,
                   max_epoch: int = 100
                   ):
    filed_dict = md.build_field()

    me_model = MulEncNet(
        target_name=md.target_key,
        text_names=md.keys,
        filed_dict=filed_dict,
        max_out=max_out,
        num_layers=layers,
        embedding_size=embedding_size,
        hidden_size=hidden_size,
        dropout=0.1,
        lr=0.001
    )

    early_stop_callback = EarlyStopping(monitor="val_bleu4",
                                        min_delta=0.00,
                                        patience=patience,
                                        verbose=False,
                                        mode="max")

    checkpoint_callback = ModelCheckpoint(save_top_k=1,
                                          monitor="val_bleu4",
                                          mode="max",
                                          filename="MulEnc-{epoch:02d}-{val_bleu4:.4f}")
    trainer = Trainer(check_val_every_n_epoch=1,
                      gpus=num_gpu,
                      gradient_clip_val=5,
                      max_epochs=max_epoch,
                      callbacks=[early_stop_callback, checkpoint_callback])

    train_set, valid_set, test_set = md.to_datasets()

    train_loader = train_set.gen_data_loader(batch_size=batch_size)
    valid_loader = valid_set.gen_data_loader(batch_size=batch_size)
    test_loader = test_set.gen_data_loader(batch_size=batch_size)

    trainer.fit(me_model, train_dataloaders=train_loader, val_dataloaders=valid_loader)

    best_model = MulEncNet.load_from_checkpoint(checkpoint_callback.best_model_path,
                                                target_name=md.target_key,
                                                text_names=md.keys,
                                                filed_dict=filed_dict)

    trainer.test(best_model, dataloaders=test_loader)
